﻿#pragma once

#include <asio.hpp>

#ifdef WIN32
#    include <process.h>
#    include <windows.h>

#    define DIRSEPA '\\'
#else
#    include <sys/sysinfo.h>
#    include <unistd.h>

#    define DIRSEPA '/'
#endif

#include <array>
#include <atomic>
#include <chrono>
#include <cmath>
#include <functional>
#include <iostream>
#include <memory>
#include <random>
#include <string_view>

namespace leak {

namespace utils {
    // =====================================================================================================================
    // Tools
    // =====================================================================================================================
    inline uint64_t GetSecondTime() {
        return std::chrono::duration_cast<std::chrono::seconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    }

    inline uint64_t GetMillisecondTime() {
        return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    }

    inline uint64_t GetMicrosecondTime() {
        return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    }

    inline uint64_t GetNanosecondTime() {
        return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    }

    inline int64_t GetRandomNumber(int64_t min, int64_t max) {
        static std::random_device              rd;
        static std::mt19937_64                 rng(rd());  // Generates random 64-bit integers
        std::uniform_int_distribution<int64_t> dist{ min, max };
        return dist(rng);
    }

    inline float GetRandomFloat(float min, float max) {
        static std::random_device             rd;
        static std::mt19937_64                rng(rd());
        std::uniform_real_distribution<float> dist{ min, max };
        return dist(rng);
    }

    inline int GetPid() {
#ifndef __linux__
        return (int)_getpid();
#else
        return (int)getpid();
#endif
    }

    inline size_t GetProcessorsNumber() {
#ifndef __linux__
        SYSTEM_INFO sysInfo;
        GetSystemInfo(&sysInfo);
        return sysInfo.dwNumberOfProcessors;
#else
        return get_nprocs_conf();
#endif
    }

    inline std::string GetApplicationDirectory() {
        std::string full_name;
        char        path[1024] = { 0 };
#ifdef _WIN32
        LPTSTR full_path = (LPTSTR)&path;
        if (::GetModuleFileName(NULL, full_path, 1024)) {
            full_name.assign(path);
        }
        std::replace(full_name.begin(), full_name.end(), DIRSEPA, '/');
#else
        int count = readlink("/proc/self/exe", path, 1024);
        if (count > 0 && count < 1024) {
            path[count] = '\0';
            full_name.assign(path);
        } else {
            char  cmdline[1024] = { 0 };
            FILE* fp            = NULL;

            sprintf(cmdline, "ls -lt /proc/%d | grep exe | awk '{print $NF}'", getpid());
            if ((fp = popen(cmdline, "r"))) {
                if (fgets(path, 1024, fp)) {
                    size_t len = strlen(path);
                    if (len > 0) {
                        if ('\n' == path[--len])
                            path[len] = '\0';
                        full_name.assign(path);
                    }
                }
                pclose(fp);
            }
        }
#endif

        return full_name.substr(0, full_name.find_last_of("/"));
    }

    // =====================================================================================================================
    // Context pool
    // =====================================================================================================================
    class IContextPool {
    public:
        virtual ~IContextPool()          = default;
        virtual asio::io_context& Next() = 0;
    };

    // =====================================================================================================================
    // Sequence context pool
    // =====================================================================================================================
    class SequenceContextPool final : public IContextPool {
    public:
        SequenceContextPool(size_t len = 0);
        SequenceContextPool(const SequenceContextPool&)            = delete;
        SequenceContextPool& operator=(const SequenceContextPool&) = delete;
        ~SequenceContextPool();

        asio::io_context& Next() override;

    private:
        std::vector<std::thread>       threads_;
        std::vector<asio::io_context*> contexts_;
        std::atomic<uint16_t>          index_;
    };

    // =====================================================================================================================
    // Buffer
    // =====================================================================================================================
    class Buffer {
    public:
        Buffer(int init_size = 1024);
        Buffer(const Buffer&)            = delete;
        Buffer& operator=(const Buffer&) = delete;
        Buffer(Buffer&& buffer);
        ~Buffer();

        int32_t Append(std::string_view data);
        int32_t Append(const void* data, int32_t size);
        int32_t Append(const char* data, int32_t size);
        void    Retrieve();
        void    Reserve(int32_t size);

        int32_t          capacity() const { return size_; }
        int32_t          writeable_length() const { return size_ - write_; }
        int32_t          readable_length() const { return write_ - read_; }
        std::string_view read_view() const { return std::string_view(buffer_ + read_, readable_length()); }
        char*            write_addr() { return buffer_; }
        void             has_write(int32_t size) { write_ += size; }
        void             has_read(int32_t size) { read_ += size; }

    private:
        char*   buffer_;
        int32_t write_;
        int32_t read_;
        int32_t size_;
    };

    inline SequenceContextPool::SequenceContextPool(size_t len) {
        if (len == 0) len = GetProcessorsNumber();
        contexts_.reserve(len);
        threads_.reserve(len);
        for (int i = 0; i < len; i++) {
            auto context = contexts_.emplace_back(new asio::io_context);
            threads_.emplace_back([context]() {
                auto worker = asio::make_work_guard(*context);
                context->run();
            });
        }
    }

    inline SequenceContextPool::~SequenceContextPool() {
        for (auto& context : contexts_) {
            context->stop();
            delete context;
        }
        for (auto& thread : threads_) {
            thread.join();
        }
    }

    inline asio::io_context& SequenceContextPool::Next() {
        return *contexts_[index_.fetch_add(1) % contexts_.size()];
    }

    inline Buffer::Buffer(int init_size /*= 1024*/)
      : buffer_(nullptr)
      , write_(0)
      , read_(0)
      , size_(init_size) {
        buffer_ = static_cast<char*>(std::malloc(init_size));
    }

    inline Buffer::Buffer(Buffer&& buffer) {
        buffer_        = buffer.buffer_;
        write_         = buffer.write_;
        read_          = buffer.read_;
        buffer.buffer_ = nullptr;
    }

    inline Buffer::~Buffer() {
        if (buffer_) std::free(buffer_);
    }

    inline int32_t Buffer::Append(std::string_view data) {
        return Append(data.data(), data.length());
    }

    inline int32_t Buffer::Append(const void* data, int32_t size) {
        return Append(static_cast<const char*>(data), size);
    }

    inline int32_t Buffer::Append(const char* data, int32_t size) {
        Reserve(size);
        std::memcpy(buffer_ + write_, data, size);
        write_ += size;
        return size;
    }

    inline void Buffer::Retrieve() {
        if (read_ == 0) return;
        int32_t len = readable_length();
        assert(len >= 0);
        if (len) std::memmove(buffer_, buffer_ + read_, len);
        read_  = 0;
        write_ = len;
    }

    inline void Buffer::Reserve(int32_t size) {
        int need = size - writeable_length();
        if (need > 0) {
            buffer_ = static_cast<char*>(std::realloc(buffer_, this->size_ + need));
            assert(buffer_);
            this->size_ += need;
        }
    }

    using ContextPool = std::shared_ptr<IContextPool>;
}  // namespace utils

namespace net::tcp {
    template <typename T>
    struct SocketUtils {};

    template <>
    struct SocketUtils<asio::ip::tcp::socket> {
        using Socket = asio::ip::tcp::socket;

        static asio::ip::tcp::endpoint endpoint(Socket& s) { return s.remote_endpoint(); }

        static void shutdown(Socket& s) { s.shutdown(Socket::shutdown_send); }

        static void close(Socket& s) {
            asio::error_code ec;
            s.close(ec);
        }

        static decltype(auto) basic(Socket& s) { return s; }
    };

    // =====================================================================================================================
    // Socket
    // =====================================================================================================================
    template <typename Basic>
    class SocketWrapper final : public std::enable_shared_from_this<SocketWrapper<Basic>> {
    public:
        SocketWrapper(const SocketWrapper&)            = delete;
        SocketWrapper& operator=(const SocketWrapper&) = delete;
        ~SocketWrapper()                               = default;

        static std::shared_ptr<SocketWrapper<Basic>> New(Basic basic) {
            return std::shared_ptr<SocketWrapper<Basic>>(new SocketWrapper<Basic>(std::move(basic)));
        }

        void StartRead(std::function<size_t(std::shared_ptr<SocketWrapper<Basic>>, std::string_view)> oncomplete) {
            if (!open_) return;
            auto self = this->shared_from_this();
            if (recvbuffer_.writeable_length() == 0)
                recvbuffer_.Reserve(recvbuffer_.capacity());
            basic_.async_read_some(
                asio::buffer(recvbuffer_.write_addr(), recvbuffer_.writeable_length()),
                [this, self, oncomplete = std::move(oncomplete)](auto err, size_t rlen) {
                    if (!open_) return;
                    if (!err) {
                        recvbuffer_.has_write(rlen);
                        if (oncomplete)
                            rlen = std::invoke(oncomplete, self, recvbuffer_.read_view());
                        else
                            rlen = -1;
                        recvbuffer_.has_read(std::min<size_t>(rlen, recvbuffer_.readable_length()));
                        recvbuffer_.Retrieve();
                        StartRead(std::move(oncomplete));
                    } else {
                        OnError("[read] " + err.message());
                    }
                });
        }

        template <bool shutdown = false, typename... Args>
        SocketWrapper<Basic>* Write(Args&&... args) {
            if (!open_ || shutdown_send_) return this;
            if (shutdown) shutdown_send_ = true;
            auto self = this->shared_from_this();
            asio::async_write(
                basic_,
                std::array<asio::const_buffer, sizeof...(Args)>{ asio::buffer(std::forward<Args>(args))... },
                [this, self](auto err, size_t wlen) {
                    if (err) {
                        OnError("[write] " + err.message());
                    } else {
                        if (shutdown_send_) SocketUtils<Basic>::shutdown(basic_);
                    }
                });
            return this;
        }

        char* Prepare(size_t size) {
            if (!open_) return nullptr;
            sendbuffer_.Reserve(size);
            return sendbuffer_.write_addr();
        }

        SocketWrapper<Basic>* PrepareWrite(std::string_view data) {
            if (!open_ || shutdown_send_) return this;
            sendbuffer_.Append(data);
            return this;
        }

        SocketWrapper<Basic>* Flush(bool shutdown = false) {
            static constexpr size_t MAX_FLUSH_SIZE = 10 * 1024;
            if (!open_ || writing_) return this;
            if (shutdown) shutdown_send_ = true;
            if (sendbuffer_.readable_length() == 0) {
                if (shutdown_send_) SocketUtils<Basic>::shutdown(basic_);
                return this;
            }
            writing_  = true;
            auto self = this->shared_from_this();
            asio::async_write(
                basic_,
                asio::buffer(sendbuffer_.read_view().substr(0, MAX_FLUSH_SIZE)),
                [this, self](auto err, size_t wlen) {
                    if (!err) {
                        sendbuffer_.has_read(wlen);
                        sendbuffer_.Retrieve();
                        writing_ = false;
                        Flush();
                    } else {
                        OnError("[write] " + err.message());
                    }
                });
            return this;
        }

        void close() {
            SocketUtils<Basic>::close(basic_);
        }

        asio::ip::tcp::endpoint endpoint() {
            return SocketUtils<Basic>::endpoint(basic_);
        }

        decltype(auto) basic() { return SocketUtils<Basic>::basic(basic_); }

        template <typename T>
        T* data() { return std::any_cast<T>(&data_); }

        template <typename T>
        SocketWrapper<Basic>* set_data(T&& t) {
            data_ = std::forward<T>(t);
            return this;
        }

        SocketWrapper<Basic>* set_onerror(std::function<void(std::string_view)> onerror) {
            onerror_ = std::move(onerror);
            return this;
        }

        SocketWrapper<Basic>* set_open() {
            open_ = true;
            return this;
        }

        SocketWrapper<Basic>* set_no_delay(bool enable) {
            basic_.set_option(asio::ip::tcp::no_delay(enable));
            return this;
        }

        SocketWrapper<Basic>* set_reuse_address(bool enable) {
            basic_.set_option(asio::socket_base::reuse_address(enable));
            return this;
        }

    private:
        SocketWrapper(Basic basic)
          : basic_(std::move(basic))
          , writing_(false)
          , open_(false)
          , shutdown_send_(false) {
        }

        void OnError(std::string_view err) {
            if (!open_) return;
            SocketUtils<Basic>::close(basic_);
            recvbuffer_.Retrieve();
            sendbuffer_.Retrieve();
            open_          = false;
            writing_       = false;
            shutdown_send_ = false;
            if (onerror_) std::invoke(onerror_, err);
        }

        Basic         basic_;
        utils::Buffer recvbuffer_;
        utils::Buffer sendbuffer_;
        bool          writing_;
        bool          open_;
        bool          shutdown_send_;

        std::any                              data_;
        std::function<void(std::string_view)> onerror_;
    };

    using Socket = SocketWrapper<asio::ip::tcp::socket>;

    // =====================================================================================================================
    // Server
    // =====================================================================================================================
    class Server : public std::enable_shared_from_this<Server> {
        using ErrorHandler = std::function<void(std::string_view)>;

    public:
        ~Server() { Close(); }

        static std::shared_ptr<Server> New(asio::io_context& context) {
            return std::shared_ptr<Server>(new Server(context));
        }

        Server* Listen(std::string_view host, uint16_t port) {
            if (acceptor_.is_open()) return this;

            asio::ip::tcp::endpoint ep(asio::ip::make_address(host), port);
            acceptor_.open(ep.protocol());
            acceptor_.set_option(asio::socket_base::reuse_address(true));
            acceptor_.bind(ep);
            acceptor_.listen();
            return this;
        }

        void Close() {
            if (acceptor_.is_open()) {
                acceptor_.close();
            }
        }

        void Accept(std::function<void(std::shared_ptr<Socket>)> onaccept) {
            asio::io_context& context = context_pool_ ? context_pool_->Next() : context_;
            auto              socket  = Socket::New(asio::ip::tcp::socket(context));
            auto              self    = shared_from_this();
            acceptor_.async_accept(socket->basic(), [self, socket, this, onaccept = std::move(onaccept)](auto err) {
                if (!err) {
                    socket->set_open()
                        ->set_no_delay(true)
                        ->set_reuse_address(true);
                    if (onaccept) std::invoke(onaccept, socket);
                    Accept(std::move(onaccept));
                } else if (onaccept_error_) {
                    std::invoke(onaccept_error_, err.message());
                }
            });
        }

        Server* set_context_pool(utils::ContextPool pool) {
            context_pool_ = pool;
            return this;
        }
        Server* set_onaccept_error(ErrorHandler fn) {
            onaccept_error_ = std::move(fn);
            return this;
        }
        asio::ip::tcp::endpoint endpoint() {
            return acceptor_.local_endpoint();
        }

    private:
        Server(asio::io_context& context)
          : context_(context)
          , acceptor_(context) {
        }

        asio::io_context&       context_;
        asio::ip::tcp::acceptor acceptor_;
        utils::ContextPool      context_pool_;
        ErrorHandler            onaccept_error_;
    };

    // =====================================================================================================================
    // Connector
    // =====================================================================================================================
    template <typename C, typename E>
    inline decltype(auto) Connect(asio::io_context& context, std::string_view host, uint16_t port, C onconnect, E onerror) {
        auto resolver = std::make_shared<asio::ip::tcp::resolver>(context);
        auto socket   = Socket::New(asio::ip::tcp::socket(context));
        resolver->async_resolve(
            host.data(),
            std::to_string(port),
            [resolver, socket, onconnect = std::move(onconnect), onerror = std::move(onerror)](auto ec, auto results) {
                if (!ec) {
                    asio::async_connect(socket->basic(), results, [socket, onconnect = std::move(onconnect)](auto ec, auto ep) {
                        if (!ec) {
                            socket->set_open()
                                ->set_no_delay(true)
                                ->set_reuse_address(true);
                            std::invoke(onconnect, socket);
                        }
                    });
                } else {
                    std::invoke(onerror, ec.message());
                }
            });
        return socket;
    }

    template <typename S, typename C, typename E>
    inline void Connect(S socket, std::string_view host, uint16_t port, C onconnect, E onerror) {
        auto resolver = std::make_shared<asio::ip::tcp::resolver>(socket->basic().get_executor());
        resolver->async_resolve(
            host.data(),
            std::to_string(port),
            [resolver, socket, onconnect = std::move(onconnect), onerror = std::move(onerror)](auto ec, auto results) {
                if (!ec) {
                    asio::async_connect(socket->basic(), results, [socket, onconnect = std::move(onconnect)](auto ec, auto ep) {
                        if (!ec) {
                            socket->set_open()
                                ->set_no_delay(true)
                                ->set_reuse_address(true);
                            std::invoke(onconnect, socket);
                        }
                    });
                } else {
                    std::invoke(onerror, ec.message());
                }
            });
    }
}  // namespace net::tcp

}  // namespace leak